import datetime
from collections import Counter

from bot.openai_bot import OpenAIBot
from enums.nwes_type_enum import NewsTypeEnum
from items.response import GenericResponse
from items.task import Task
from service.baike_title_service import BaikeTitleService
from service.decoration_title_service import DecorationTitleService
from service.strategy_title_service import StrategyTitleService
from utils.constants import TAGS, DECORATION_TAGS, STRATEGY_TAGS


def create_classification_task(
        task: Task,
):
    """
    创建分类任务
    :param task: 待分类的任务对象
    :return: 任务id
    """
    # 将标题向量化
    openai_bot = OpenAIBot()
    vector = openai_bot.embeddings(text_list=[task.title, ])[0]
    res_ls = list()
    # 新增分类类型判断
    if task.type == NewsTypeEnum.WIKI.value:
        baike_title_service = BaikeTitleService()
        res_ls = baike_title_service.get_title_tags(title_vector=vector, top_k=10)
    elif task.type == NewsTypeEnum.DECORATION.value:
        decoration_title_service = DecorationTitleService()
        res_ls = decoration_title_service.get_title_tags(title_vector=vector, top_k=10)
    elif task.type == NewsTypeEnum.STRATEGY.value:
        strategy_title_service = StrategyTitleService()
        res_ls = strategy_title_service.get_title_tags(title_vector=vector, top_k=10)
    # 结果解析
    # 1.详情
    predict_result_details = list()
    tags = list()
    tag_type = {}
    if task.type == NewsTypeEnum.WIKI.value:
        tag_type = TAGS.copy()
    elif task.type == NewsTypeEnum.DECORATION.value:
        tag_type = DECORATION_TAGS.copy()
    elif task.type == NewsTypeEnum.STRATEGY.value:
        tag_type = STRATEGY_TAGS.copy()
    for index, item in enumerate(res_ls):
        title = item.get("title")
        title_tag = tag_type.get(item.get("title_tag"))
        score = round(item.get("distance")*100, 2)
        predict_result_details.append({
            "title": title,
            "title_tag": title_tag,
            "score": score
        })
        tags.append(title_tag)

    # 2.最高分
    tag_max_score = tag_type.get(res_ls[0].get("title_tag"))

    # 3.最高频
    tag_counts = Counter(tags)
    tag_most_common = tag_counts.most_common(n=1)[0][0]

    return GenericResponse(
        now=int(datetime.datetime.now().timestamp()),
        data={
            "result": {
                "predict_result": tag_max_score,
                "predict_result_most_common": tag_most_common,
                "predict_result_details": predict_result_details,
                "classifier": "vector_retrieval",
            }
        }
    )


def main():
    task = Task()
    task.title = "房产税"
    res = create_classification_task(task)
    print(res)


if __name__ == '__main__':
    main()
